#include <optix.h>
#include <optixu/optixu_math_namespace.h>
#include "device_util.h"
#include "random.h" 

using namespace optix;

rtDeclareVariable(rtObject,      top_object, , );
rtDeclareVariable(float3, shading_normal, attribute shading_normal, ); 
rtDeclareVariable(float3, geometric_normal, attribute geometric_normal, ); 
rtDeclareVariable(float,      t_hit,        rtIntersectionDistance, );
rtDeclareVariable(optix::Ray, ray,          rtCurrentRay, );

// Max number of photon bounces
rtDeclareVariable(uint, max_photon_depth, , );

rtDeclareVariable(uint2, launch_index, rtLaunchIndex, );
rtBuffer<float4, 2>   result_buffer;
rtBuffer<ViewRayData, 2>  hit_buffer;
rtBuffer<uint2, 2>		seed_buffer;
rtBuffer<PhotonRec, 1> photon_buffer;
rtDeclareVariable(uint, maxPhotonsPerRay, , );

rtDeclareVariable(PhotonRayData, ray_payload, rtPayload, );
rtDeclareVariable(uint2, launch_dim,   rtLaunchDim, );

rtDeclareVariable(float3,        eye, , );
rtDeclareVariable(float3,        U, , );
rtDeclareVariable(float3,        V, , );
rtDeclareVariable(float3,        W, , );

rtDeclareVariable(Light,				 light , , );

rtDeclareVariable(float3,  Ks, , );
rtDeclareVariable(float3,  Kd, , );

RT_PROGRAM void closest_hit()
{
	// Check if this is a light source
  float3 world_shading_normal   = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, shading_normal ) );
  float3 world_geometric_normal = normalize( rtTransformNormal( RT_OBJECT_TO_WORLD, geometric_normal ) );
  float3 ffnormal     = faceforward( world_shading_normal, -ray.direction, world_geometric_normal );

  float3 hit_point = ray.origin + t_hit*ray.direction;
  float3 newDir;
	 

  //if( fmaxf( Kd ) > 0.0f ) {
    // We hit a diffuse surface; record hit if it has bounced at least once
   if( ray_payload.depth >= 0 ) {
      PhotonRec& rec = photon_buffer[ launch_index.x * launch_dim.y * maxPhotonsPerRay + launch_index.y * maxPhotonsPerRay + ray_payload.nPhotons];
      rec.pos = hit_point;
      rec.normal = ffnormal;
      rec.incDir = ray.direction;
      rec.energy = ray_payload.energy;
      ray_payload.nPhotons++;
    }

    ray_payload.energy = Kd * ray_payload.energy; 
    float3 U,V,W;
    createONB(ffnormal, U, V, W);
		float2 rnd2 = make_float2( rnd(ray_payload.sample.x), rnd(ray_payload.sample.y) );
    sampleUnitHemisphere( rnd2, U, V, W, newDir );			
  /*
	} else {
    ray_payload.energy = Ks * ray_payload.energy;
    // Find reflection dir
    newDir = reflect( ray.direction, ffnormal );
  }
	*/

  
  if ( ray_payload.nPhotons >= maxPhotonsPerRay || ray_payload.depth >= max_photon_depth) { 
    return;
	}

	ray_payload.depth++;
  optix::Ray new_ray( hit_point, newDir, 1, 1e-4 );
  rtTrace(top_object, new_ray, ray_payload);

}


__device__ __inline__ void generateAreaLightPhoton( const Light& light, const float2& d_sample, float3& o, float3& d)
{
  // Choose a random position on light
  o = light.anchor + 0.5f * ( light.v1 + light.v2 );
  
  // Choose a random direction from light
  float3 U, V, W;
  createONB( light.direction, U, V, W);
  sampleUnitHemisphere( d_sample, U, V, W, d );
}

RT_PROGRAM void cast() 
{
	float2 screen = make_float2( seed_buffer.size() ); 
	float2 sample2;
	float3 U,V,W,dir,origin;
	uint2 seed = seed_buffer[launch_index];
	sample2 = make_float2(
      ( static_cast<float>( launch_index.x ) + rnd( seed.x ) ) / static_cast<float>( screen.x ),
      ( static_cast<float>( launch_index.y ) + rnd( seed.y ) ) / static_cast<float>( screen.y ) );
	
	
	createONB(light.direction, U, V, W);
	sampleUnitHemisphere( sample2, U, V, W, dir );
	
	origin = light.anchor + 0.5f*(light.v1 + light.v2); //rnd( seed.x ) * light.v1 + rnd( seed.y ) * light.v2;
	dir = normalize(make_float3(-1.f, -1.f, 0.f));

	float side = 400.f;
	float3 target = make_float3( -side/2 + rnd(seed.x) * side, 0.0f, -side/2 + rnd(seed.y) * side );

	//dir = normalize(target - origin);
	//generateAreaLightPhoton( light, sample2, origin, dir );
	
	optix::Ray ray( origin, dir, 1, 1e-4f );

	PhotonRayData prd;
  //  rec.ray_dir = ray_direction; // set in ppass_closest_hit
  prd.energy = light.power;
  prd.sample = seed; 
  prd.nPhotons = 0;
  prd.depth = 0;

	PhotonRec& rec = photon_buffer[ launch_index.x * launch_dim.y * maxPhotonsPerRay + launch_index.y * maxPhotonsPerRay ];
  rec.energy = make_float3(-1.f); // last element flag

	rtTrace( top_object, ray, prd );
}

RT_PROGRAM void except()
{
  const unsigned int code = rtGetExceptionCode();
  rtPrintf( "Caught photon exception 0x%X at launch index (%d,%d)\n", code, launch_index.x, launch_index.y );

	result_buffer[launch_index] = make_float4(1.f, 0.f, 1.f, 0.f);
}